/*************************************************************************
				FFTSHARE.C

	Title:		FFT-preparing functions for Speden.
	Author:		Hanna Szoke
	Date:		Oct 24, 2005 
	Function:       This package contains initializations, 
                        convolve, deconvolve and do_fft for Solve and
			Back, using NDFT, as well as code for Regrid
			using FFTW3 (equal-to-equal spaced grids).

*************************************************************************/

#include "util.h"
#include "dims.h"
#include "sp.h"
#include <nfft.h>		 
#include <fftw3.h>		 

nfft_plan	plan ;
static	fftw_complex	*f ;
static	fftw_complex	*f_hat ;

extern	real  aAxis, bAxis, cAxis ;	/* dimensions of unit cell. */

#define MIN(a,b) (a<b?a:b)

void initfft(int type, int N)    /* N is Nfc, Nfob or their sum.  */
{
     int my_N[3],my_n[3];
     int	n, i, p ;
     int	sfc = sizeof(fftw_complex) ;
   
     /* initialise the plan */

     if  (N < 1.e24) {
        printTwice("Using original NDFT initialization parameters\n") ;
        my_N[0]=Nx; my_n[0]=2*next_power_of_2(Nx);
        my_N[1]=Ny; my_n[1]=2*next_power_of_2(Ny);
        my_N[2]=Nz; my_n[2]=2*next_power_of_2(Nz);
        p = MIN(Nz,MIN(Ny,MIN(3,Nx))) ;
     }
     else {
        printTwice("Using reduced NDFT initialization parameters\n") ;
        my_N[0]=Nx; my_n[0]=2*(Nx);
        my_N[1]=Ny; my_n[1]=2*(Ny);
        my_N[2]=Nz; my_n[2]=2*(Nz);
        p = MIN(Nz,MIN(Ny,MIN(2,Nx))) ;
     }
 
     f = (fftw_complex *) e_malloc(N*sfc, "initfft") ;
     f_hat = (fftw_complex *) e_malloc(Nptotal*sfc, "initfft") ;

/* prepare coordinates on non-equidistant grid.
   they are re-scaled to span (-1/2 to 1/2)   */

     nfft_init_specific(&plan, 3, my_N, N, my_n, p,
                      PRE_PHI_HUT| PRE_PSI| MALLOC_X |
                      FFTW_INIT| FFT_OUT_OF_PLACE,
                      FFTW_MEASURE| FFTW_DESTROY_INPUT);

     if (type == 1) {			/* Back */
        for (n = 0; n < Nfc; n++) {
	     *(plan.x + 3*n)     = Fc_Hlist[n] / ((aAxis)) ;
	     *(plan.x + 3*n + 1) = Fc_Klist[n] / ((bAxis)) ;
	     *(plan.x + 3*n + 2) = Fc_Llist[n] / ((cAxis)) ;
	}
     }
     else if (type == 2){ 		/* Solve */
        for (n = 0; n < Nfob; n++) {
	     *(plan.x + 3*n)     = Hlist[n] / ((aAxis)) ;
	     *(plan.x + 3*n + 1) = Klist[n] / ((bAxis)) ;
	     *(plan.x + 3*n + 2) = Llist[n] / ((cAxis)) ;
        }
     }
     else {				/* solve with phase extension */
        for (n = 0; n < Nfob; n++) {
	     *(plan.x + 3*n)     = Hlist[n] / ((aAxis)) ;
	     *(plan.x + 3*n + 1) = Klist[n] / ((bAxis)) ;
	     *(plan.x + 3*n + 2) = Llist[n] / ((cAxis)) ;
        }
        for (n = Nfob, i = 0; n < N; n++, i++) {
	     *(plan.x + 3*n)     = Fc_Hlist[i] / ((aAxis)) ;
	     *(plan.x + 3*n + 1) = Fc_Klist[i] / ((bAxis)) ;
	     *(plan.x + 3*n + 2) = Fc_Llist[i] / ((cAxis)) ;
        }
     }

     if( plan.nfft_flags & PRE_PSI )
	nfft_precompute_psi( &plan );
} 

void	convolve(int N, real *array, real *efac, COMPLEX *prod)
{
	/***********************************************************
	 Accumulate complex product of DFT(array) and (real) efac.

	 On entry:
		array = Current value of the image.
	        efac = real array to multiply Fourier transform
	 On return:
	 	prod = complex product
	***********************************************************/

        int	i ;

      	plan.f_hat = f_hat;
        plan.f = f;

        for (i = 0; i < Nptotal; i++) {
        	plan.f_hat[i][0] = array[i];
                plan.f_hat[i][1] = 0;     
        }
        nfft_trafo(&plan) ;

        for (i = 0; i < N; i++)
        {
   	     prod[i].re = efac[i] * f[i][0] ;
	     prod[i].im = efac[i] * f[i][1];
        }
}
  


void	deconvolve(int N, real *garray, real *efac, COMPLEX *kern)
{
	/***********************************************************
	 Calculate the gradient from kern and ADD it to the possibly
	 non-zero garray.
	
	 On entry:
	       efac = real array to multiply Fourier transform
	       kern = r, weighted & with exp_phi
	 On return:
		garray  += gradient 
	***********************************************************/

	int     i ;

        for (i = 0; i < N; i++)
	{
		f[i][0] = efac[i] * kern[i].re ;
		f[i][1] = efac[i] * kern[i].im ;
        }

        plan.f_hat = f_hat;
        plan.f = f;

        nfft_adjoint(&plan) ;

        for (i = 0; i < Nptotal; i++)
	     garray[i] += -plan.f_hat[i][0] ;
}

void	do_fft(COMPLEX *myarray, int isign)
{
     int	n;

	/***************************************************************
	Following usual convention, isign is -1 for forward, 1 for back 
	transform.  
	***************************************************************/

      if (isign == -1) {	/* equidistant to non-e	*/
      	   plan.f_hat = f_hat;
	   plan.f = f;
	   for (n = 0; n < Nptotal; n++) {
		   plan.f_hat[n][0] = myarray[n].re ;
		   plan.f_hat[n][1] = myarray[n].im ;
           }
           nfft_trafo(&plan) ;
      }
      else {
      	   plan.f_hat = f_hat;
	   plan.f = f;
           nfft_adjoint(&plan) ;

           for (n = 0; n < Nptotal; n++) {
               myarray[n].re = plan.f_hat[n][0] ;
               myarray[n].im = plan.f_hat[n][1] ;
           }
     }
}

/*  From here on, code is for Regrid   */

static  fftw_plan     fplan, bplan ;	

void reg_initfft()
{
     fftw_complex *embed, *trans;

     embed = (fftw_complex*)fftw_malloc( Np*sizeof( fftw_complex ) );
     trans = (fftw_complex*)fftw_malloc( Np*sizeof( fftw_complex ) );

     fplan = fftw_plan_dft_3d( Nz, Ny, Nx, embed, trans, 
                              FFTW_FORWARD, FFTW_MEASURE );
     if( !fplan )
	Fatal("Problem creating FFTW forward plan.");

     bplan = fftw_plan_dft_3d(Nz, Ny, Nx, embed, trans, 
                              FFTW_BACKWARD, FFTW_MEASURE) ;
     if (!bplan)
          Fatal("Problem creating FFTW backward plan") ;
}
void    reg_do_fft(COMPLEX *inarray, int isign)
{
        /***************************************************************
        Following usual convention, isign is -1 for forward, 1 for back
        transform.  
        ***************************************************************/

      if (isign == -1)
           fftw_execute(fplan) ;
      else
           fftw_execute(bplan) ;
}
